import asyncio
from .program_db import ProgramDB
from .llm_client import LLMClient
from .prompt_sampler import PromptSampler
from .evaluator import Evaluator
from .utils import apply_diff
import random

class AlphaEvolve:
    """
    协调进化过程的主类。
    """
    
    def __init__(self, program_db: ProgramDB, llm_client: LLMClient,
                 prompt_sampler: PromptSampler, evaluator: Evaluator, num_iterations: int,elite_selection_pressure: float = 0.1):
        self.program_db = program_db
        self.llm_client = llm_client
        self.prompt_sampler = prompt_sampler
        self.evaluator = evaluator
        self.num_iterations = num_iterations
        self.elite_selection_pressure = elite_selection_pressure

    
    async def run(self):
        """
        运行指定迭代次数的进化循环，并提供详细的日志。
        """
        for i in range(self.num_iterations):
            print(f"\n{'=' * 25} 迭代 {i + 1}/{self.num_iterations} {'=' * 25}")
            
            if not self.program_db.population:
                print("[ERROR] 种群为空，进化过程无法继续。")
                break
            
            # 1. SELECTION - 选择父代-改进的选择策略
            if random.random() < self.elite_selection_pressure:
                # 利用：选择最好的程序
                parent_program, parent_score_dict = self.program_db.get_best_program()
                selection_method = "精英选择 (Exploitation)"
            else:
                # 探索：随机选择
                parent_program, parent_score_dict = self.program_db.get_random_program_with_score()
                selection_method = "随机选择 (Exploration)"

            parent_score = parent_score_dict['main_score']
            print(f"[1. SELECTION] 使用 '{selection_method}' 选定父代程序 (当前分数: {parent_score})")
            
            # 2. PROMPTING & GENERATION - 提示构建与生成
            prompt = self.prompt_sampler.build_prompt(parent_program, self.program_db.get_best_programs())
            print("[2. GENERATION] 正在向 LLM 请求代码变更...")
            diff = await self.llm_client.generate_code_diff(prompt)
            
            if not diff or "SEARCH" not in diff or "REPLACE" not in diff:
                print("[GENERATION] LLM 未能生成有效的 Diff，跳过此轮。")
                continue
            
            # 打印 LLM 生成的 Diff，这是可观测性的关键
            print("[GENERATION] LLM 生成的 Diff 如下:")
            print("----------------- DIFF START -----------------")
            print(diff)
            print("------------------ DIFF END ------------------")
            
            # 3. MUTATION - 突变
            try:
                child_program = apply_diff(parent_program, diff)
                print("[3. MUTATION] 成功应用 Diff，生成子代程序。")
            except ValueError as e:
                print(f"[MUTATION] 应用 Diff 失败: {e}，跳过此轮。")
                continue
            
            # 4. EVALUATION - 评估
            print("[4. EVALUATION] 正在评估子代程序...")
            score = await self.evaluator.evaluate(child_program)
            
            # 5. SURVIVAL - 存活
            if score and score.get("main_score") is not None:
                new_score = score['main_score']
                
                # --- 添加对比日志 ---
                comparison_text = ""
                if self.program_db._is_better(new_score, parent_score):
                    comparison_text = " (优于父代)"
                elif new_score == parent_score:
                    comparison_text = " (与父代相同)"
                else:
                    comparison_text = " (差于父代)"
                # ---------------------
                
                print(f"[5. SURVIVAL] 子代评估完成. 新分数: {new_score}{comparison_text}")
                self.program_db.add_program(child_program, score)
                _, best_score_dict = self.program_db.get_best_program()
                print(f"[5. SURVIVAL] 新的种群最佳分数: {best_score_dict['main_score']}")
            
            # --- 每轮结束时打印种群状态 ---
            print("\n--- 当前种群状态 (Top 3) ---")
            top_programs = self.program_db.get_best_programs(n=3)
            for idx, (_, score_dict) in enumerate(top_programs):
                print(f"  #{idx + 1}: 分数 = {score_dict['main_score']}")
            print("=" * 62)
        
        return self.program_db.get_best_program()